import os
import pandas as pd
import requests
import json

# -----------初始化表格-----------
# 初始化表格，提取特定列
def initialize_table():
    columns_to_keep = ['攻击者IP', '受害者IP', '攻击手段', '攻击次数']
    for file in os.listdir('.'):
        if file.endswith('.xlsx'):
            df = pd.read_excel(file)
            new_df = df[columns_to_keep]
            try:
                new_df.to_excel('data.xlsx', index=False)
                print("初始化成功......")
                return
            except Exception as e:
                print(f"保存Excel文件data.xlsx时出现错误: {e}")
                raise
    print("当前目录下未找到.xlsx格式的表格文件")


# -----------分类表格-----------
# 定义判断IP是否为私网IP的函数
def is_private_ip(ip):
    ip_parts = ip.split('.')
    if len(ip_parts) != 4:  # 先判断IP格式是否正确，若不正确则返回False（可按需调整处理方式）
        return False
    first_octet = int(ip_parts[0])
    second_octet = int(ip_parts[1]) if len(ip_parts) > 1 else None
    if first_octet == 10:
        return True
    elif first_octet == 172 and 16 <= second_octet <= 31:
        return True
    elif first_octet == 192 and ip_parts[1] == '168':
        return True
    return False


# -----------分类公网内网IP-----------
# 分类公网内网IP
def sort_ips():
    df_data = pd.read_excel("data.xlsx")
    df_data['IP类型'] = df_data['受害者IP'].apply(is_private_ip).map({True: "私网IP", False: "公网IP"})
    private_df = df_data[df_data['IP类型'] == "私网IP"].drop('IP类型', axis=1)
    public_df = df_data[df_data['IP类型'] == "公网IP"].drop('IP类型', axis=1)
    private_df.to_excel("政务区.xlsx", index=False)
    print("政务区.xlsx创建成功")
    public_df.to_excel("公众区.xlsx", index=False)
    print("公众区.xlsx创建成功")


# -----------获取归属地-----------
# 调取api获取归属地
def get_ip_info(file_name):
    file_city = pd.read_excel(file_name)
    cols = file_city.columns.tolist()
    ip_index = cols.index('攻击者IP')
    count_index = cols.index('攻击次数')
    file_city['归属地'] = ""
    cols.insert(ip_index + 1, '归属地')
    file_city = file_city[cols]
    for index, row in file_city.iterrows():
        ip = row['攻击者IP']
        url = f'https://api.vvhan.com/api/ipInfo?ip={ip}'
        try:
            response = requests.get(url)
            dic = json.loads(response.text)
            ip_info = dic.get('info', {}).get('city', '未知')
            if ip_info == '-':
                ip_info = dic.get('info', {}).get('prov', '未知')
                if ip_info == '-':
                    ip_info = dic.get('info', {}).get('country', '未知')
            file_city.at[index, '归属地'] = ip_info
        except:
            file_city.at[index, '归属地'] = '获取失败'
    file_city.to_excel(file_name, index=False, engine='openpyxl')
    print(file_name + "创建成功")


# -----------分类汇总-----------
# 分类汇总函数定义
def group_by(xlsx_name, roles):
    xlsx_file = pd.read_excel(xlsx_name)
    xlsx_name = xlsx_name.split('.')

    gp = xlsx_file.groupby(roles)['攻击次数'].sum().sort_values(ascending=False)
    result = pd.DataFrame(gp).reset_index()
    result.columns = [roles, '攻击次数']
    new_name = xlsx_name[0] + roles + "汇总.xlsx"
    result.to_excel(new_name, index=False)
    if new_name == "公众区攻击者IP汇总.xlsx":
        get_ip_info(new_name)
    else:
        print(new_name + "创建成功")


if __name__ == "__main__":
    initialize_table()
    sort_ips()
    group_by("公众区.xlsx", "攻击者IP")
    group_by("公众区.xlsx", "受害者IP")
    group_by("公众区.xlsx", "攻击手段")

    group_by("政务区.xlsx", "攻击者IP")
    group_by("政务区.xlsx", "受害者IP")
    group_by("政务区.xlsx", "攻击手段")